/***************************************************************************
 *   Copyright (C) 2010 by programmeboy                                    *
 *   programmeboy@gmail.com                                                *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.                                        *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             * 
 ***************************************************************************/

#include <ntddk.h>
#include "vt.h"
#include "cpu.h"

#define kmalloc(_s)	ExAllocatePoolWithTag(NonPagedPool, _s, 'neek')
#define kfree(_p)	ExFreePool(_p)

typedef struct _MULTI_CPU_CONTEXT
{
    ULONG NumberOfRaisedCPU;
    ULONG AllCPURaised;
    PKDPC BasePKDPC;
    KIRQL SavedIRQL;

}MULTI_CPU_CONTEXT, *PMULTI_CPU_CONTEXT;

static MULTI_CPU_CONTEXT MultiCPUContext = {0};

VOID RaiseCPUIrqlAndWait(
     IN PKDPC Dpc,
     IN PVOID DeferredContext,
     IN PVOID SystemArgument1,
     IN PVOID SystemArgument2
    )
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    InterlockedIncrement( &MultiCPUContext.NumberOfRaisedCPU );
	while ( !InterlockedCompareExchange( 
        &MultiCPUContext.AllCPURaised, 1, 1 ) ){

		DO_NOTHING;
	}
	InterlockedDecrement( &MultiCPUContext.NumberOfRaisedCPU );
}

BOOLEAN 
GainExlusivity(
        VOID
        )
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	
	NTSTATUS ntStatus;
	ULONG CurrentCPU;
	CCHAR i;
	PKDPC tmpPKDPC;
	
	KIRQL NewIrql;

	NewIrql = KeRaiseIrqlToDpcLevel();
	
	//
	//  Initialize the various
	//
	
	InterlockedAnd(&MultiCPUContext.NumberOfRaisedCPU, 0);
	InterlockedAnd(&MultiCPUContext.AllCPURaised, 0);
	
	//
	//  Allocate KDPC for each Processor
	//
	
	tmpPKDPC = (PKDPC)kmalloc( KeNumberProcessors * sizeof(KDPC) );
	
	if (NULL == tmpPKDPC){
		KfLowerIrql(NewIrql);
		return FALSE;
	}
	
	MultiCPUContext.BasePKDPC = tmpPKDPC;
	CurrentCPU = KeGetCurrentProcessorNumber();
	
	for (i = 0; i < KeNumberProcessors; i++, *tmpPKDPC++){
        
        //
        //  Set the loop DPC for the cpu that not current's
        //
        
		if ( i != CurrentCPU ){
			KeInitializeDpc( tmpPKDPC, RaiseCPUIrqlAndWait, NULL );
			KeSetTargetProcessorDpc( tmpPKDPC, i );
			KeInsertQueueDpc( tmpPKDPC, NULL, NULL );
		}
	}
    
    //
    //  Wait for cpu enter loop.
    //
    
	while( (KeNumberProcessors - 1) != 
        InterlockedCompareExchange( &MultiCPUContext.NumberOfRaisedCPU, \
        KeNumberProcessors - 1, KeNumberProcessors - 1) ){
		DO_NOTHING;
	}
    
    //
    //  Save the new IRQL
    //
    
    MultiCPUContext.SavedIRQL = NewIrql;
	return TRUE;
	
}

VOID 
ReleaseExclusivity()
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    InterlockedIncrement(&MultiCPUContext.AllCPURaised);

	while (InterlockedCompareExchange( 
           &MultiCPUContext.NumberOfRaisedCPU, 0, 0)){
		DO_NOTHING;
	}
    
    //
    //  Free the KDPC
    //
    
	if (NULL != MultiCPUContext.BasePKDPC){
		kfree((PVOID)MultiCPUContext.BasePKDPC);
		MultiCPUContext.BasePKDPC = NULL;
	}
	
	KfLowerIrql(MultiCPUContext.SavedIRQL);	
}

VOID 
SetWpAndReleaseCPU()
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    _EnableWP();
    ReleaseExclusivity();
}

BOOLEAN 
ClearWpAndGainCPU()
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
	BOOLEAN bRet = GainExlusivity();
    _DisableWP();
	return bRet;
}